import numpy as np
import pdb


class Girk_curve():
    '''
    Girmmann problem's curve boundary
    '''
    def __init__(self):
        self.rho_0 = 15.0
        self.theta = 2*np.pi/9 #对应于alpha
        self.r_0 = self.rho_0/np.sin(self.theta)
        self.d = 0.06

    def phi(self,p,xi):
        #p.shape = (NEbd,2,2)
        #xi.shape = (NQ,) or (NQ,NEbd)

        #pdb.set_trace()
        s=self.Get_theta(p) #(NEbd,2)

        
        idx, = np.nonzero((np.mean(s,axis=-1)>1e-15)&(np.mean(s,axis=-1)<self.theta-1e-15)) #只有(0,theta)的需要做曲边变换


        if len(xi.shape)>1:
            shape = xi.shape[:-1] + p.shape[:-2] + (2,) #(NQ,NEbd,2)
            s = np.einsum('...i,i->...i',xi,s[...,0])+np.einsum('...i,i->...i',1-xi,s[...,1]) #(NQ,NEbd)

        else:
            shape = xi.shape + p.shape[:-2] + (2,) #(NQ,NEbd,2)
            s = np.einsum('...,i->...i',xi,s[...,0])+np.einsum('...,i->...i',1-xi,s[...,1]) #(NQ,NEbd)

        

        pp = np.zeros(shape,dtype=p.dtype)

        
        if len(idx) > 0:
            pp[...,idx,0] = np.sin(s[...,idx])
            pp[...,idx,1] = np.cos(s[...,idx])

        

        r = np.mean(self.Get_r(p),axis=-1) #(NEbd,)


        idx, = np.nonzero(np.abs(r-(self.r_0-self.d/2))<1e-15)

        pp[...,idx,:] = pp[...,idx,:]*(self.r_0-self.d/2)



        idx, = np.nonzero(np.abs(r-(self.r_0+self.d/2))<1e-15)

        pp[...,idx,:] = pp[...,idx,:]*(self.r_0+self.d/2)

        return pp




    def D_xi_phi(self,p,xi):
        #p.shape = (NEbd,2,2)
        #xi.shape = (NQ,) or (NQ,NEbd)

        s=self.Get_theta(p) #(NEbd,2)

        idx, = np.nonzero((np.mean(s,axis=-1)>1e-15)&(np.mean(s,axis=-1)<self.theta-1e-15)) #只有(0,theta)的需要做曲边变换


        if len(xi.shape)>1:
            shape = xi.shape[:-1] + p.shape[:-2] + (2,) #(NQ,NEbd,2)
            s_xi = np.einsum('...i,i->...i',xi,s[...,0])+np.einsum('...i,i->...i',1-xi,s[...,1]) #(NQ,NEbd)

        else:
            shape = xi.shape + p.shape[:-2] + (2,) #(NQ,NEbd,2)
            s_xi = np.einsum('...,i->...i',xi,s[...,0])+np.einsum('...,i->...i',1-xi,s[...,1]) #(NQ,NEbd)



        pp = np.zeros(shape,dtype=p.dtype)


        if len(idx) > 0:
            pp[...,idx,0] = np.cos(s_xi[...,idx])*(s[...,idx,0]-s[...,idx,1])
            pp[...,idx,1] = -np.sin(s_xi[...,idx])*(s[...,idx,0]-s[...,idx,1])



        r = np.mean(self.Get_r(p),axis=-1) #(NEbd,)


        idx, = np.nonzero(np.abs(r-(self.r_0-self.d/2))<1e-15)

        pp[...,idx,:] = pp[...,idx,:]*(self.r_0-self.d/2)


        idx, = np.nonzero(np.abs(r-(self.r_0+self.d/2))<1e-15)

        pp[...,idx,:] = pp[...,idx,:]*(self.r_0+self.d/2)

        return pp





















    
    def Get_theta(self,p):
        '''
        对应于alpha
        '''

        shape = p.shape[:-1]
        p = p.reshape(-1,2) #(N,2)
        r = self.Get_r(p)#(N,)
        s = np.arccos(p[:,1]/r) #[0,pi]

        if (np.sum(s>np.pi/2)>0):
            raise ValueError('some points are not right!')

        return s.reshape(shape)
    
    def Get_r(self,p):
        '''
        得到极坐标r
        '''
        return np.sqrt(np.sum(p**2,axis=-1))
